import logging
from pathlib import Path
from typing import Optional

from langchain_aws import ChatBedrock
from langchain_core.messages import SystemMessage, HumanMessage

from app.config import settings

logger = logging.getLogger(__name__)


async def optimize_with_bedrock_llm(file_path: Path, language: Optional[str] = None) -> str:
    """
    Optimize markdown document using AWS Bedrock LLM.
    
    Args:
        file_path: Path to the markdown file to optimize
        language: Target language for optimization
        
    Returns:
        Optimized document content as string
        
    Raises:
        Exception: If LLM optimization fails
    """
    content = file_path.read_text(encoding="utf-8")
    prompt = settings.system_prompt.format(language)
    user_content = settings.user_content.format(content)

    # Initialize Bedrock LLM with credentials if provided
    bedrock_kwargs = {
        "model_id": settings.bedrock_model,
        "model_kwargs": {
            "max_tokens": 16384,
        },
        "streaming": True,
        "region_name": settings.bedrock_region,
    }
    
    # Add credentials if provided
    if settings.bedrock_access_key_id and settings.bedrock_secret_access_key:
        bedrock_kwargs["aws_access_key_id"] = settings.bedrock_access_key_id
        bedrock_kwargs["aws_secret_access_key"] = settings.bedrock_secret_access_key
    
    llm = ChatBedrock(**bedrock_kwargs)

    messages = [
        SystemMessage(content=prompt),
        HumanMessage(content=user_content),
    ]

    try:
        # Use streaming
        result = []
        async for chunk in llm.astream(messages):
            if hasattr(chunk, "content") and chunk.content:
                result.append(str(chunk.content))
        return "".join(result)
    except Exception as e:
        logger.error(f"Bedrock LLM optimization failed for {file_path}: {str(e)}")
        # Re-raise the exception so it can be caught by the calling function
        raise


async def describe_image_with_bedrock_llm(image_path: Path, language: Optional[str] = None) -> str:
    """
    Use AWS Bedrock LLM to describe the content of an image, with a prompt in English suitable for RAG documentation.
    
    Args:
        image_path: Path to the image file to describe
        language: Target language for description
        
    Returns:
        Image description as string
        
    Raises:
        Exception: If LLM image description fails
    """
    from app.utils.image_utils import image_to_base64
    base64_str = image_to_base64(image_path)

    # English prompt for RAG
    prompt = (
        f"You are an expert in enterprise knowledge base documentation. Based on the provided image, directly output a detailed description of the image content, without any formatting instructions. "
        f"The description should be concise and accurate, suitable for insertion into a Retrieval-Augmented Generation (RAG) document, and optimized for downstream knowledge retrieval and question answering. "
        f"Language: {language or 'English'}."
    )
    user_content = f"The base64 content of the image is: data:image/png;base64,{base64_str}"

    # Initialize Bedrock LLM with credentials if provided
    bedrock_kwargs = {
        "model_id": settings.bedrock_model,
        "model_kwargs": {
            "max_tokens": 16384,
        },
        "streaming": True,
        "region_name": settings.bedrock_region,
    }
    
    # Add credentials if provided
    if settings.bedrock_access_key_id and settings.bedrock_secret_access_key:
        bedrock_kwargs["aws_access_key_id"] = settings.bedrock_access_key_id
        bedrock_kwargs["aws_secret_access_key"] = settings.bedrock_secret_access_key
    
    llm = ChatBedrock(**bedrock_kwargs)

    messages = [
        SystemMessage(content=prompt),
        HumanMessage(content=user_content),
    ]

    try:
        result = []
        async for chunk in llm.astream(messages):
            if hasattr(chunk, "content") and chunk.content:
                result.append(str(chunk.content))
        return "".join(result)
    except Exception as e:
        logger.error(f"Bedrock LLM image description failed for {image_path}: {str(e)}")
        # Re-raise the exception so it can be caught by the calling function
        raise
